import os
import random
import logging
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class InpaintingDataset(Dataset):
    """用于图像修复任务的数据集"""
    def __init__(self, root_dir, mode='train', image_size=256, data_percentage=100):
        """
        初始化数据集
        
        参数:
            root_dir (str): 数据集根目录
            mode (str): 'train' 或 'val'，用于确定使用的数据集部分
            image_size (int): 图像调整大小的目标尺寸
            data_percentage (int): 使用的数据百分比（1-100）
        """
        self.root_dir = root_dir
        self.mode = mode
        self.image_size = image_size
        self.data_percentage = data_percentage
        
        # 根据模式设置不同的路径结构
        if mode == 'train':
            self.mask_dir = os.path.join(root_dir, 'train', 'mask')
            self.gt_dir = os.path.join(root_dir, 'train', 'gt')
        elif mode == 'val':
            self.mask_dir = os.path.join(root_dir, 'test_val', 'val_256', 'mask')
            self.gt_dir = os.path.join(root_dir, 'test_val', 'val_256', 'gt')
        else:
            raise ValueError(f"不支持的模式: {mode}")
        
        # 构建图像路径对列表
        self.image_pairs = []
        self._build_image_pairs()
        
        # 为RGB和灰度图像创建单独的变换
        self.rgb_transform = transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])  # RGB图像
        ])
        
        self.mask_transform = transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor()  # 掩码不需要归一化
        ])
        
    def _build_image_pairs(self):
        """构建图像路径对列表，考虑训练集和验证集的不同路径结构"""
        if self.mode == 'train':
            # 训练集路径模式: /root/code/datasets/Places365_256x256/train/mask/a/airfield/a/airfield/00000001.jpg
            #               /root/code/datasets/Places365_256x256/train/gt/a/airfield/00000001.jpg
            for alphabet_dir in os.listdir(self.gt_dir):
                alphabet_path = os.path.join(self.gt_dir, alphabet_dir)
                if os.path.isdir(alphabet_path):
                    for category_dir in os.listdir(alphabet_path):
                        category_path = os.path.join(alphabet_path, category_dir)
                        if os.path.isdir(category_path):
                            for img_name in os.listdir(category_path):
                                if img_name.endswith(('.jpg', '.png')):
                                    # 构建原始图像路径
                                    gt_path = os.path.join(category_path, img_name)
                                    
                                    # 构建对应的掩码图像路径
                                    mask_path = os.path.join(
                                        self.mask_dir, 
                                        alphabet_dir,
                                        category_dir, 
                                        alphabet_dir,
                                        category_dir, 
                                        img_name
                                    )
                                    
                                    # 确保两个文件都存在
                                    if os.path.exists(gt_path) and os.path.exists(mask_path):
                                        self.image_pairs.append((mask_path, gt_path))
        
        elif self.mode == 'val':
            # 验证集路径模式: /root/code/datasets/Places365_256x256/test_val/val_256/mask/Places365_val_00000001.jpg
            #               /root/code/datasets/Places365_256x256/test_val/val_256/gt/Places365_val_00000001.jpg
            for img_name in os.listdir(self.gt_dir):
                if img_name.endswith(('.jpg', '.png')):
                    # 构建原始图像路径
                    gt_path = os.path.join(self.gt_dir, img_name)
                    
                    # 构建对应的掩码图像路径
                    mask_path = os.path.join(self.mask_dir, img_name)
                    
                    # 确保两个文件都存在
                    if os.path.exists(gt_path) and os.path.exists(mask_path):
                        self.image_pairs.append((mask_path, gt_path))
        
        # 根据设定的百分比随机选择数据
        total_images = len(self.image_pairs)
        if self.data_percentage < 100:
            num_to_keep = int(total_images * self.data_percentage / 100)
            random.shuffle(self.image_pairs)
            self.image_pairs = self.image_pairs[:num_to_keep]
        
        logger.info(f"模式: {self.mode}, 找到 {total_images} 对图像, 使用 {len(self.image_pairs)} 对 ({self.data_percentage}%)")
    
    def __len__(self):
        return len(self.image_pairs)
    
    def __getitem__(self, idx):
        mask_path, gt_path = self.image_pairs[idx]
        
        try:
            # 读取掩码和原始图像
            mask_img = Image.open(mask_path).convert('L')  # 灰度图像
            gt_img = Image.open(gt_path).convert('RGB')
            
            # 应用变换
            mask_img = self.mask_transform(mask_img)
            gt_img = self.rgb_transform(gt_img)
            
            # 将掩码二值化: 1表示需要修复区域，0表示保持原样
            mask = (mask_img > 0.5).float()
            
            # 创建带有"洞"的图像
            masked_img = gt_img * (1 - mask)
            
            return {
                'masked_image': masked_img,  # 带有缺失区域的图像
                'mask': mask,                # 掩码（1表示需要修复区域）
                'gt_image': gt_img           # 原始图像（目标）
            }
        except Exception as e:
            logger.error(f"处理图像时出错: {mask_path}, {gt_path}, 错误: {str(e)}")
            # 返回一个替代项以防止训练中断
            return self.__getitem__((idx + 1) % len(self.image_pairs))